package org.finesys.chat.core.base.agent;


import java.lang.reflect.Method;
import java.lang.reflect.Parameter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

import org.finesys.ai.core.chat.json.JsonObjectSchema;
import org.finesys.ai.core.chat.json.JsonSchemaElement;
import org.finesys.chat.core.base.json.JsonSchemaElementHelper;
import org.finesys.common.core.util.Utils;

public class ToolSpecifications {

    private ToolSpecifications() {
    }

    public static List<ToolSpecification> toolSpecificationsFrom(Class<?> classWithTools) {
        List<ToolSpecification> toolSpecifications = (List<ToolSpecification>) Arrays.stream(classWithTools.getDeclaredMethods()).filter((method) -> {
            return method.isAnnotationPresent(Tool.class);
        }).map(ToolSpecifications::toolSpecificationFrom).collect(Collectors.toList());
        validateSpecifications(toolSpecifications);
        return toolSpecifications;
    }

    public static List<ToolSpecification> toolSpecificationsFrom(Object objectWithTools) {
        return toolSpecificationsFrom(objectWithTools.getClass());
    }

    public static void validateSpecifications(List<ToolSpecification> toolSpecifications) throws IllegalArgumentException {
        Set<String> names = new HashSet<>();
        Iterator<ToolSpecification> iterator = toolSpecifications.iterator();

        ToolSpecification toolSpecification;
        do {
            if (!iterator.hasNext()) {
                return;
            }

            toolSpecification = (ToolSpecification) iterator.next();
        } while (names.add(toolSpecification.name()));

        throw new IllegalArgumentException(String.format("Tool names must be unique. The tool '%s' appears several times", toolSpecification.name()));
    }

    public static ToolSpecification toolSpecificationFrom(Method method) {
        Tool annotation = (Tool) method.getAnnotation(Tool.class);
        String name = Utils.isNullOrBlank(annotation.name()) ? method.getName() : annotation.name();
        String description = String.join("\n", annotation.value());
        if (description.isEmpty()) {
            description = null;
        }

        JsonObjectSchema parameters = parametersFrom(method.getParameters());
        return ToolSpecification.builder().name(name).description(description).parameters(parameters).build();
    }

    private static JsonObjectSchema parametersFrom(Parameter[] parameters) {
        Map<String, JsonSchemaElement> properties = new LinkedHashMap<>();
        List<String> required = new ArrayList<>();
        Map<Class<?>, JsonSchemaElementHelper.VisitedClassMetadata> visited = new LinkedHashMap<>();
        int size = parameters.length;

        for (int i = 0; i < size; ++i) {
            Parameter parameter = parameters[i];
            if (!parameter.isAnnotationPresent(ToolMemoryId.class)) {
                boolean isRequired = (Boolean) Optional.ofNullable((P) parameter.getAnnotation(P.class)).map(P::required).orElse(true);
                properties.put(parameter.getName(), jsonSchemaElementFrom(parameter, visited));
                if (isRequired) {
                    required.add(parameter.getName());
                }
            }
        }

        Map<String, JsonSchemaElement> definitions = new LinkedHashMap<>();
        visited.forEach((clazz, visitedClassMetadata) -> {
            if (visitedClassMetadata.recursionDetected) {
                definitions.put(visitedClassMetadata.reference, visitedClassMetadata.jsonSchemaElement);
            }

        });
        if (properties.isEmpty()) {
            return null;
        } else {
            return JsonObjectSchema.builder().properties(properties).required(required).definitions(definitions.isEmpty() ? null : definitions).build();
        }
    }

    private static JsonSchemaElement jsonSchemaElementFrom(Parameter parameter, Map<Class<?>, JsonSchemaElementHelper.VisitedClassMetadata> visited) {
        P annotation = (P) parameter.getAnnotation(P.class);
        String description = annotation == null ? null : annotation.value();
        return JsonSchemaElementHelper.jsonSchemaElementFrom(parameter.getType(), parameter.getParameterizedType(), description, visited);
    }
}
